import logging
from typing import List, Dict, Any
from langchain_milvus import Milvus
from langchain_ollama import OllamaEmbeddings, OllamaLLM
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate
from langchain_core.documents import Document
from pymilvus import Collection, connections, utility
from app.config import settings
from app.schemas.civil_exam_schemas import PositionInfo, UserInfo

# 初始化日志
logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)


class CivilExamRAGService:
    """公考选岗RAG服务"""

    def __init__(
            self,
            milvus_host: str = settings.MILVUS_HOST,
            milvus_port: str = settings.MILVUS_PORT,
            milvus_token: str = None,
            collection_name: str = "civil_exam_positions",
            embedding_model: str = settings.EMBEDDING_MODEL,
            llm_model: str = settings.LLM_MODEL,
            ollama_base_url: str = settings.OLLAMA_BASE_URL,
            dim: int = 1024
    ):
        """初始化公考选岗RAG服务"""
        self.milvus_host = milvus_host
        self.milvus_port = milvus_port
        self.milvus_token = milvus_token
        self.collection_name = collection_name
        self.embedding_model = embedding_model
        self.llm_model = llm_model
        self.ollama_base_url = ollama_base_url
        self.dim = dim

        # 初始化核心组件
        self.embeddings = self._init_embeddings()
        self.vector_store = self._init_vector_store()
        self.llm = self._init_llm()
        self.qa_chain = self._init_qa_chain()

    def _init_embeddings(self):
        """初始化嵌入模型"""
        try:
            embeddings = OllamaEmbeddings(
                model=self.embedding_model,
                base_url=self.ollama_base_url
            )
            logger.info(f"嵌入模型就绪：{self.embedding_model}")
            return embeddings
        except Exception as e:
            logger.error(f"嵌入模型初始化失败：{str(e)}")
            raise

    def _init_vector_store(self):
        """初始化向量存储"""
        try:
            # 连接Milvus
            connections.connect(
                host=self.milvus_host,
                port=self.milvus_port,
                token=self.milvus_token
            )

            # 初始化Milvus向量库
            vector_store = Milvus(
                embedding_function=self.embeddings,
                connection_args={
                    "host": self.milvus_host,
                    "port": self.milvus_port,
                    "token": self.milvus_token
                },
                collection_name=self.collection_name,
                drop_old=False
            )

            logger.info(f"Milvus向量库就绪：{self.collection_name}")
            return vector_store
        except Exception as e:
            logger.error(f"Milvus初始化失败：{str(e)}")
            raise

    def _init_llm(self):
        """初始化大语言模型"""
        try:
            llm = OllamaLLM(
                model=self.llm_model,
                base_url=self.ollama_base_url,
                temperature=0.1
            )
            logger.info(f"LLM就绪：{self.llm_model}")
            return llm
        except Exception as e:
            logger.error(f"LLM初始化失败：{str(e)}")
            raise

    def _init_qa_chain(self):
        """初始化QA链"""
        position_recommendation_prompt = PromptTemplate(
            template="""
任务：根据用户个人信息和公务员岗位信息，推荐最适合的岗位并给出推荐理由。

岗位信息（从向量库中检索）：{context}

用户信息：{question}

要求：
1. 分析用户信息与岗位的匹配度
2. 综合考虑专业、学历、地域偏好、工作经验等因素
3. 输出推荐的岗位列表，包含岗位ID、部门、地区、职位名称
4. 为每个推荐岗位提供推荐理由
5. 按匹配度从高到低排序
6. 仅输出JSON格式，不包含其他文字

输出格式示例：
{{
  "recommendations": [
    {{
      "position_id": "岗位ID",
      "department": "部门名称",
      "location": "地区",
      "position_name": "职位名称",
      "match_reason": "匹配理由"
    }}
  ]
}}
""",
            input_variables=["context", "question"]
        )

        return RetrievalQA.from_chain_type(
            llm=self.llm,
            chain_type="stuff",
            retriever=self.vector_store.as_retriever(search_kwargs={"k": 5}),
            chain_type_kwargs={"prompt": position_recommendation_prompt},
            return_source_documents=True
        )

    def recommend_positions(self, user_info: UserInfo) -> Dict[str, Any]:
        """根据用户信息推荐岗位"""
        try:
            # 构造查询问题
            question = f"""
用户信息：
- 学历：{user_info.education}
- 专业：{user_info.major}
- 地域偏好：{', '.join(user_info.location_preference)}
- 工作经验：{user_info.experience_years}年
- 持有证书：{', '.join(user_info.professional_certificates) if user_info.professional_certificates else '无'}
- 技能：{', '.join(user_info.skills) if user_info.skills else '无'}
"""

            # 执行查询
            result = self.qa_chain.invoke({"query": question})

            # 解析结果
            import json
            import re

            json_match = re.search(r'\{[\s\S]*?\}', result["result"].strip())
            if json_match:
                recommendations = json.loads(json_match.group())
            else:
                recommendations = {"recommendations": []}

            return {
                "code": 200,
                "message": "推荐成功",
                "data": recommendations,
                "source_documents": [
                    {
                        "content": doc.page_content,
                        "metadata": doc.metadata
                    }
                    for doc in result["source_documents"]
                ]
            }
        except Exception as e:
            logger.error(f"岗位推荐失败：{str(e)}")
            return {
                "code": 500,
                "message": f"岗位推荐失败：{str(e)}",
                "data": {}
            }

    def add_position_info(self, position_info: PositionInfo):
        """添加岗位信息到向量库"""
        try:
            # 将岗位信息转换为文档格式
            content = f"""
岗位ID：{position_info.position_id}
部门：{position_info.department}
地区：{position_info.location}
职位名称：{position_info.position_name}
学历要求：{position_info.education_requirement}
专业要求：{position_info.major_requirement}
工作经验要求：{position_info.experience_requirement or '无要求'}
薪资范围：{position_info.salary_range or '面议'}
招聘人数：{position_info.recruitment_number}
考试科目：{', '.join(position_info.exam_subjects)}
职位描述：{position_info.job_description}
职位要求：{position_info.job_requirements}
额外福利：{position_info.extra_benefits or '无'}
"""

            metadata = {
                "position_id": position_info.position_id,
                "department": position_info.department,
                "location": position_info.location,
                "position_name": position_info.position_name,
                "education_requirement": position_info.education_requirement,
                "major_requirement": position_info.major_requirement
            }

            # 创建文档对象
            doc = Document(page_content=content, metadata=metadata)

            # 添加到向量库
            self.vector_store.add_documents([doc])

            logger.info(f"岗位信息添加成功：{position_info.position_id}")
            return True
        except Exception as e:
            logger.error(f"添加岗位信息失败：{str(e)}")
            return False